from transformers import BertTokenizer,BertForSequenceClassification,pipeline

model_dir = r"model/bert-base-chinese"

# 加载模型和分词器，本地没有回直接下载
model = BertForSequenceClassification.from_pretrained("bert-base-chinese", cache_dir = model_dir)
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese", cache_dir = model_dir)

# 创建分类pipeline
classifier = pipeline("text-classification", model = model, tokenizer = tokenizer, device = "cuda")

# 进行文本分类
result = classifier("你好，我是一款语言模型")
print(result)
